package cn.da.shuai.cool.mybatis.interceptor;

import cn.da.shuai.cool.mybatis.annotation.DataScope;
import cn.da.shuai.cool.mybatis.enums.DataScopeType;
import cn.da.shuai.cool.mybatis.service.ICoolDataAuthService;
import cn.hutool.core.collection.ListUtil;
import cn.hutool.core.util.ArrayUtil;
import cn.hutool.core.util.StrUtil;
import net.sf.jsqlparser.expression.Expression;
import net.sf.jsqlparser.expression.operators.conditional.AndExpression;
import net.sf.jsqlparser.parser.CCJSqlParserUtil;
import net.sf.jsqlparser.statement.select.PlainSelect;
import org.apache.ibatis.executor.statement.StatementHandler;
import org.apache.ibatis.mapping.BoundSql;
import org.apache.ibatis.mapping.MappedStatement;
import org.apache.ibatis.plugin.Interceptor;
import org.apache.ibatis.plugin.Intercepts;
import org.apache.ibatis.plugin.Invocation;
import org.apache.ibatis.plugin.Signature;
import org.apache.ibatis.reflection.MetaObject;
import org.apache.ibatis.reflection.SystemMetaObject;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

import java.lang.reflect.Method;
import java.sql.Connection;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;

/**
 * 数据权限拦截器
 */
@Component
@Intercepts({@Signature(type = StatementHandler.class, method = "prepare", args = {Connection.class, Integer.class})})
public class DataScopeInterceptor implements Interceptor {

    @Autowired(required = false)
    private ICoolDataAuthService dataAuthService;

    @Override
    public Object intercept(Invocation invocation) throws Throwable {

        if (dataAuthService == null) {
            return invocation.proceed();
        }

        StatementHandler statementHandler = (StatementHandler) invocation.getTarget();
        MetaObject metaObject = SystemMetaObject.forObject(statementHandler);
        MappedStatement mappedStatement = (MappedStatement) metaObject.getValue("delegate.mappedStatement");

        // 获取执行的方法
        String id = mappedStatement.getId();
        String className = id.substring(0, id.lastIndexOf("."));
        String methodName = id.substring(id.lastIndexOf(".") + 1);
        Class<?> clazz = Class.forName(className);
        Method[] methods = clazz.getDeclaredMethods();

        DataScope dataScope = null;
        for (Method method : methods) {
            if (method.getName().equals(methodName)) {
                if (method.isAnnotationPresent(DataScope.class)) {
                    dataScope = method.getAnnotation(DataScope.class);
                    break;
                }
            }
        }

        // 如果没有数据权限注解，则不处理
        if (dataScope == null) {
            return invocation.proceed();
        }

        BoundSql boundSql = (BoundSql) metaObject.getValue("delegate.boundSql");
        String originalSql = boundSql.getSql();

        // 解析SQL
        PlainSelect plainSelect = (PlainSelect) CCJSqlParserUtil.parse(originalSql);

        // 构建数据权限的SQL条件
        Expression whereExpression = plainSelect.getWhere();
        String tableAlias = dataScope.tableAlias();
        String prefix = tableAlias.isEmpty() ? "" : tableAlias + ".";

        Expression dataFilterExpression;
        DataScopeType[] types = dataScope.types();
        if (ArrayUtil.isEmpty(types)) {
            dataFilterExpression = CCJSqlParserUtil.parseCondExpression("1 = 0");
        } else {
            // 1.如果types中包含ALL，直接通过
            if (Arrays.asList(types).contains(DataScopeType.ALL)) {
                return invocation.proceed();
            }

            // 2.通过接口获取到所有用户的id列表
            List<Integer> userIds = dataAuthService.getUserIdsByDataScope(types);

            // 3.如果id过多，需要进行拆分
            List<List<Integer>> userIdSplits = ListUtil.split(userIds, 1000);
            if (userIdSplits.size() == 1) {
                dataFilterExpression = CCJSqlParserUtil.parseCondExpression(
                        prefix + dataScope.userIdColumn() + " in ( " + StrUtil.join(",", userIdSplits.getFirst()) + ") ");
            } else {
                // 处理拆分的数据
                List<Expression> conditions = new ArrayList<>();
                for (List<Integer> userIdList : userIdSplits) {
                    conditions.add(CCJSqlParserUtil.parseCondExpression(
                            prefix + dataScope.userIdColumn() + " in ( " + StrUtil.join(",", userIdList) + ") "));
                }

                dataFilterExpression = conditions.getFirst();
                for (int i=1; i<conditions.size(); i++) {
                    dataFilterExpression = new AndExpression(dataFilterExpression, conditions.get(i));
                }
            }
        }

        if (whereExpression == null) {
            plainSelect.setWhere(dataFilterExpression);
        } else {
            plainSelect.setWhere(new AndExpression(whereExpression, dataFilterExpression));
        }

        // 替换原始SQL
        metaObject.setValue("delegate.boundSql.sql", plainSelect.toString());

        return invocation.proceed();
    }

}